/*
 * Copyright (C) 2023 Coder.AN
 * Email: an.hongjun@foxmail.com
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <https://www.gnu.org/licenses/>.
 */
#ifndef __EDGE_TENSOR_H__
#define __EDGE_TENSOR_H__


#include <vector>
#include <iostream>
#include <stdint.h>
#include <cuda_runtime.h>


#define CHECK(status) \
    do\
    {\
        auto ret = (status);\
        if (ret != 0)\
        {\
            std::cerr << "Cuda failure: " << ret << "at" << __FILE__ << ":" << __LINE__ << std::endl;\
            abort();\
        }\
    } while (0)


namespace tensor
{

struct Dim
{
    size_t N, C, W, H;
    Dim()
    {
        N = C = W = H = 0;
    }
    Dim(int n, int c, int w, int h)
    {
        N = n;
        C = c;
        W = w;
        H = h;
    }
};

struct Point
{
    int w, h;
    Point()
    {
        w = h = 0;
    }
    Point(int ww, int hh)
    {
        w = ww;
        h = hh;
    }
};

struct Color
{
    uint8_t x, y, z;
    Color()
    {
        x = y = z;
    }
    Color(uint8_t xx, uint8_t yy, uint8_t zz)
    {
        x = xx;
        y = yy;
        z = zz;
    }
};

enum DimOrder
{
    err_order = -1,
    NHWC = 0,
    NCWH = 1,
    NCHW = 2,
    WNCH = 3,
};

enum DType
{
    err_type = -1,
    uint8   = 0,
    uint16  = 1,
    uint32  = 2,
    uint64  = 3,
    int8    = 4,
    int16   = 5,
    int32   = 6,
    int64   = 7,
    float32 = 8,
    float64 = 9,
};

enum DColor
{
    err_color = -1,
    COLOR_RGB = 0,
    COLOR_BGR = 1,
    COLOR_GRAY = 2
};

class Tensor
{
private:
    Dim dim;
    DimOrder dim_order;
    DType dType;

    uint8_t* cpu_buffer;
    void* gpu_buffer;

    size_t unit_size;
    size_t unit_count;
    size_t total_size;

    bool empty;

public:
    Tensor();
    Tensor(void* cpu_buffer, Dim ddim, DimOrder ddim_order, DType dtype);
    Tensor(Dim ddim, DimOrder ddim_order, DType dtype);
    ~Tensor();

    // init & clear
    void update_from_cpu(void* cpu_buffer);
    void update_from_gpu(void* new_gpu_buffer);
    void set_shape(Dim ddim, DimOrder ddim_order, DType dtype);
    void clear();

    // basic algorithm
    void resize(Dim new_dim);
    void rect(Point p1, Point p2);
    void transpose(DimOrder target_order);
    void convertType(DType target_type);
    void convertColor(DColor src_color, DColor dst_color);
    void normalize();
    void hflip();
    void vflip();
    void hconcat(Tensor& t);
    void vconcat(Tensor& t);

    //draw
    void draw_line(Point start, Point end, uint8_t grayVal, int r=2);
    void draw_line(Point start, Point end, Color colorVal, int r=2);
    void draw_circle(Point center, uint8_t grayVal, int r=4);
    void draw_circle(Point center, Color colorVal, int r=4);

    // export
    void* cpu();
    void print();
    void copyTo(Tensor* t);
    
    // tools for segmentation
    void map_seg_res(std::vector<uint8_t> map_list);

    // get tensor info
    void* get_val();
    Dim get_dim();
    DType get_dtype();
    DimOrder get_order();
    bool isEmpty();
};

};

#endif // __CUDATENSOR_H__